import torch.nn as nn
import torch
from diffurec import DiffuRec
import numpy as np


class AttnDiffuRec(nn.Module):
    '''
        完整的 DiffuRec 模型
    '''

    def __init__(self, diffu, args):
        super(AttnDiffuRec, self).__init__()
        self.emb_dim = args.hidden_size
        self.item_num = args.item_num+1
        self.item_embeddings = nn.Embedding(self.item_num, self.emb_dim)
        self.embed_dropout = nn.Dropout(args.emb_dropout)
        self.position_embeddings = nn.Embedding(args.max_len, args.hidden_size)
        # LayerNorm for NLP   # 512*50*128
        self.LayerNorm = nn.LayerNorm(self.emb_dim)
        self.dropout = nn.Dropout(args.dropout)
        self.diffu = diffu
        self.loss_ce = nn.CrossEntropyLoss()
        self.loss_ce_rec = nn.CrossEntropyLoss(reduction='none')
        self.loss_mse = nn.MSELoss()

    def diffu_predict(self, item_emb, target_emb, mask_seq):
        x_0, new_item_emb,  weights, t = self.diffu(
            item_emb, target_emb, mask_seq)
        return x_0, new_item_emb, weights, t

    def forward(self, sequence: torch.Tensor, target: torch.Tensor, train_flag: bool = True):
        '''
        + sequence (batch_size, seq_len): 商品序列
        + target (batch_size,): 目标商品
        + train_flag: 真则训练，假则预测
        '''
        item_emb = self.item_embeddings(
            sequence)   # (batch_size, seq_len, emb_size)
        item_emb = self.embed_dropout(item_emb)
        item_emb = self.LayerNorm(item_emb)
        mask_seq = (sequence > 0).float()   # (batch_size,): 记录交互记录是否存在

        if train_flag:
            target_emb = self.item_embeddings(target.squeeze(-1))
            rep_diffu, rep_item, weights, t = self.diffu_predict(
                item_emb, target_emb, mask_seq)
            item_rep_dis = None         # ?
            seq_rep_dis = None

        scores = None
        return scores, rep_diffu, weights, t, item_rep_dis, seq_rep_dis
